import numpy as np
import pandas as pd
import warnings
from scipy.stats import pearsonr

from bin_quant.agent.quant_agent.factor import FactorCalculations
warnings.filterwarnings('ignore')

class FactorCalculator:
    def __init__(self, data):
        self.data = data.copy()
        self.grouped = self.data.groupby('stock')
    def calculate_factors(self):
        """
        计算所有10个因子
        :return: 包含所有计算因子的DataFrame
        """
        # 计算每个因子并添加到DataFrame中
        self.data['momentum_5'] = FactorCalculations.calculate_momentum_5(self.data, self.grouped)
        self.data['reversal_20'] = FactorCalculations.calculate_reversal_20(self.data, self.grouped)
        self.data['volatility_20'] = FactorCalculations.calculate_volatility_20(self.data, self.grouped)
        self.data['volume_change_5'] = FactorCalculations.calculate_volume_change_5(self.data, self.grouped)
        self.data['pe_ratio'] = FactorCalculations.calculate_pe_ratio(self.data, None)
        self.data['pb_ratio'] = FactorCalculations.calculate_pb_ratio(self.data, None)
        self.data['size_factor'] = FactorCalculations.calculate_size_factor(self.data, None)
        self.data['liquidity'] = FactorCalculations.calculate_liquidity(self.data, None)
        self.data['hl_volatility'] = FactorCalculations.calculate_hl_volatility(self.data, None)
        self.data['rsi_5'] = FactorCalculations.calculate_rsi_5(self.data, self.grouped)
        v = self.data

        return self.data

    @staticmethod
    def calculate_ic(factor, forward_return):
        """计算信息系数(IC)"""
        valid_idx = (~np.isnan(factor)) & (~np.isnan(forward_return))
        if sum(valid_idx) < 2:
            return np.nan
        return pearsonr(factor[valid_idx], forward_return[valid_idx])[0]

    @staticmethod
    def calculate_ir(ic_series):
        """计算信息比率(IR)"""
        return np.nanmean(ic_series) / np.nanstd(ic_series)

    def evaluate_factors(self):
        """评估因子"""
        ic_results = {}
        ic_series = {}

        factor_cols = [col for col in self.data.columns if col not in
                       ['date', 'stock', 'open', 'high', 'low', 'close', 'volume',
                        'pe', 'pb', 'market_cap', 'return']]

        for factor in factor_cols:
            ic = self.data.groupby('date').apply(
                lambda x: self.calculate_ic(x[factor], x['return']))
            ic_results[factor] = {
                'IC Mean': np.nanmean(ic),
                'IC Std': np.nanstd(ic),
                'IR': self.calculate_ir(ic),
                'IC>0 Ratio': (ic > 0).mean()
            }
            ic_series[factor] = ic

        ic_df = pd.DataFrame(ic_results).T
        return ic_df, ic_series